from run import app
import arrow
import json
from app.db.models.jobs import Job, CustomSecurityEventSettings, CustomSslSettings, CustomVulSettings
from app.db.models.tasks import Task
from app.db.models.assets import Asset, CustomTaskSettings
from app.libs import scanner
from app.libs.utils import schedule_vul, schedule_ssl, schedule_security_event, schedule_asset, schedule_http, \
    schedule_ping, schedule_setting


def cal_next_time(dt, unit, num):
    dt = arrow.Arrow.fromdatetime(dt)
    while dt < arrow.Arrow.utcnow():
        dt = dt.shift(**{unit: num})
    return dt


def get_unit_num(unit, num):
    if unit in ['month']:
        return 'days', num * 30
    if unit in ['week']:
        return 'days', num * 7
    if unit in ['day', 'd']:
        return 'days', num
    if unit in ['minute', 'm']:
        return 'minutes', num
    return unit, num


def handle_next_exec_time(setting):
    # print(setting.interval.to_dict())
    unit, num = get_unit_num(setting.interval.unit, setting.interval.num)
    if not setting.interval.startTime:
        return arrow.Arrow.utcnow()
    dt = cal_next_time(setting.interval.startTime, unit, num)
    # print(dt.for_json())
    setting.interval.startTime = dt.datetime
    # print(setting.interval.startTime)
    return dt


job_ids = []
error_jobs = []
task_ids = set()
send = False
only_exec_one_job = True

info = f"""
###############################################
开始执行周期任务迁移...
是否发送任务: {send}
是否发送一个job的任务后终止执行: {only_exec_one_job}
job_ids: {job_ids}
###############################################
"""

print(info)

if job_ids:
    jobs = Job.objects.filter(enableMonitor=True, pk__in=job_ids)
else:
    jobs = Job.objects.filter(enableMonitor=True)

i = 1
for job in jobs:
    try:
        asset = Asset.objects.with_id(job.assetId)
        print(f'{job.pk} -- {job.targetUrl}  --{i}')

        print('    schedule vul')
        handle_next_exec_time(job.vulSettings)
        if send:
            schedule_vul(job, job.vulSettings, immediate_exec=False)

        print('    schedule ssl')
        handle_next_exec_time(job.sslSettings)
        if send:
            schedule_ssl(job, job.sslSettings, immediate_exec=False)

        print('    schedule security_event')
        handle_next_exec_time(job.securityEventSettings)
        if send:
            schedule_security_event(job, job.securityEventSettings, immediate_exec=False)

        print('    schedule asset')
        handle_next_exec_time(asset.taskSettings)
        if send:
            schedule_asset(job, asset, asset.taskSettings, immediate_exec=False)

        print('    schedule ping')
        for setting in job.pingSettings:
            handle_next_exec_time(setting)
        if send:
            schedule_ping(job, job.pingSettings, immediate_exec=False)

        print('    schedule http')
        for setting in job.httpSettings:
            handle_next_exec_time(setting)
        if send:
            schedule_http(job, job.httpSettings, immediate_exec=False)
        if send:
            job.validate()
            job.save()
            asset.save()
    except:
        error_jobs.append(str(job.pk))
    finally:
        # print(job.get_task_ids(asset=asset))
        task_ids.update(job.get_task_ids(asset=asset))

    if only_exec_one_job:
        break
    # print()
    i += 1

task_id_groups = []
for task_id in task_ids:
    task_id_groups.append({
        'task_id': task_id
    })

print(f"\n删除调度任务需要的参数:\n{json.dumps(task_id_groups, indent=4)}")

print(f"\n失败的job_id:\n{error_jobs}")
